import argparse
import os
import platform
import shutil
import time
from pathlib import Path

import cv2
import torch
import torch.backends.cudnn as cudnn
from numpy import random

# https://github.com/pytorch/pytorch/issues/3678
import sys
sys.path.insert(0, './yolov5')

from yolov5.models.experimental import attempt_load
from yolov5.utils.datasets import LoadStreams, LoadImages
from yolov5.utils.general import (
    check_img_size, non_max_suppression, apply_classifier, scale_coords, xyxy2xywh, plot_one_box, strip_optimizer)
from yolov5.utils.torch_utils import select_device, load_classifier, time_synchronized

from deep_sort.utils.parser import get_config
from deep_sort.deep_sort import DeepSort

#### Add Libraries facenet_pytorch
from facenet_pytorch import MTCNN, InceptionResnetV1
from PIL import Image
import numpy as np
import mysql.connector as mydb
import json
import copy
import concurrent.futures
import torchvision.transforms as transforms

######## Global Variable

#### MTCNN ResNet のモデル読み込み
FaceSize = 60 # 切り取る最小の顔サイズ
thresholds=[0.6,0.7,0.7] # 切り出し閾値：高いほど厳しい
mtcnn = MTCNN(min_face_size=FaceSize, thresholds=thresholds, device="cuda:0")
resnet = InceptionResnetV1(pretrained='vggface2').eval()

#### データベース準備
# コネクションの作成
cnx= mydb.connect(
    host='localhost',
    #host='127.0.0.1',
    port='3306',
    user='root',
    password='dk12345678',
    database='new_db',
    #buffered = True, # https://blog.csdn.net/qq_43471020/article/details/84323154
    #connection_timeout=180,
    ssl_disabled= True
)
# カーソル作成・テーブル作成
cur = cnx.cursor()
#cur.execute("drop table faces")
cur.execute("create table if not exists faces(id integer auto_increment, feature_vector json, face_id integer, primary key (id))")
cur.execute("delete from faces")
cur.execute("truncate table faces")
#cur.execute("set transaction isolation level serializable") # トランザクションレベル変更
#cur.execute("set transaction isolation level READ UNCOMMITTED") # トランザクションレベル変更

####
simID_faceID = {}

#### IDinfoDict 
class IDInfo:
    def __init__(self, deepsort_id):
        self.d_id = deepsort_id # DeepSortによって与えられたID
        self.frame = -1 # 撮影されると毎フレームごとにプラス1
        self.Maxframe = 2 # N(=Maxframe)回に一回処理することで負荷を減らす
        self.FV_List = [] # 特徴ベクトルを保存するリスト
        self.prob_List = [] # 顔らしさを保存するリスト
        self.FV_ave = None # 特徴ベクトルの平均
        self.prob_ave = None # 顔らしさの平均
        self.certaintyID = None # この顔であると確信している顔
        self.certainty = 0 # この顔であるという確信度
        self.dissim = 0 # 誰にも似ていないフレームごとにプラス1
        self.Maxdissim = 5 # N(=Maxframe)回に一回処理することで負荷を減らす

        self.MaxFV = 5 # 特徴ベクトルを保存するリストの保存数
        self.simDict = {} # DB内の各平均特徴ベクトルと自身の平均特徴ベクトルとの類似度辞書
        self.face_id = -1 # このDeepSortIDと紐づけられているDB顔ID
        
        #self.sim


    def frameplus(self): # 毎フレーム1プラスする、N(=Maxframe)回に一回処理することで負荷を減らす
        if(self.frame == self.Maxframe): self.frame = 0
        else: self.frame += 1
    
    def addfv(self, fv, prob): # 特徴ベクトルを保存するリストに特徴ベクトルを追加する
        self.FV_List.append(fv)
        self.prob_List.append(prob)
        if len(self.FV_List) > self.MaxFV:
            del self.FV_List[0]
            del self.prob_List[0]
        self.FV_ave = np.mean(self.FV_List, axis=0)
        self.prob_ave = np.mean(self.prob_List, axis=0)

    
    def add_newID_to_DB(self): # 新しい人物のFeature Vector をデータベースに送る。生成されたIDを返す。
        if( len(self.FV_List) < self.MaxFV * 0.8): return False # データが少ない時は処理しない
        fv_json = json.dumps(self.FV_ave.tolist())
        sql = "insert into faces (feature_vector) values (\'%s\')" % fv_json
        cur.execute(sql)
        cnx.commit()
        cur.execute("select last_insert_id()")
        rows = cur.fetchall()
        new_id = rows[0][0]
        self.face_id = new_id
        if(self.certaintyID is None): simID_faceID[new_id] = new_id
        else: simID_faceID[new_id] = self.certaintyID
        return True
    
    def update_FV_to_DB(self): # 新しい人物のFeature Vector をデータベースに送る。生成されたIDを返す。
        if( len(self.FV_List) < self.MaxFV * 0.8): return False # データが少ない時は処理しない
        if( (self.prob_ave) < 0.85): return False # 顔らしくない時は処理しない
        fv_json = json.dumps(self.FV_ave.tolist())
        #sql = "insert into faces (feature_vector) values (\'%s\')" % fv_json
        sql = "update faces set feature_vector=\'%s\' where id=\'%s\'" % (fv_json, self.face_id)
        try:
            cur.execute(sql)
            cnx.commit()
        except Exception as e:
            print("ERROR",e)
            return False

        # cur.execute("select last_insert_id()")
        # rows = cur.fetchall()
        # new_id = rows[0][0]
        # self.face_id = new_id
        return True

    # 与えたデータベースrowsの各特徴ベクトルと自身の特徴ベクトルとの類似度を計算して、類似度の最大値とそのIDを返す
    def calc_sim_with_db(self, rows):
        for row in rows:
            data = json.loads(row[1])
            sim = cosine_similarity(self.FV_ave, np.array(data))
            self.simDict[row[0]] = sim
            print("sim",sim)
        MaxSim = max(self.simDict.values())
        MaxSimID = max(self.simDict, key=self.simDict.get)
        MaxSimDer = (MaxSim - np.mean(list(self.simDict.values()) ) ) / np.std(list(self.simDict.values()) )
        print("Derivation: ", MaxSimDer)
        ## 確実に本人らしい顔であればself.certaintyIDに登録
        if self.certainty < min(MaxSim, self.prob_ave):
            self.certainty = min(MaxSim, self.prob_ave)
            if(self.certainty >= 0.85):
                self.certaintyID = self.face_id
        return MaxSim, simID_faceID[MaxSimID]

### compute_color_for_labels()用　もとからあった
palette = (2 ** 11 - 1, 2 ** 15 - 1, 2 ** 20 - 1)


######## Initialize Model
out = 'inference/output'
device = select_device("0")
weights = 'yolov5/weights/yolov5m.pt'
if os.path.exists(out):
    shutil.rmtree(out)  # delete output folder
os.makedirs(out)  # make new output folder
half = device.type != 'cpu'  # half precision only supported on CUDA

# Load model
#google_utils.attempt_download(weights)
model = torch.load(weights, map_location=device)['model'].float()  # load to FP32
#model = torch.save(torch.load(weights, map_location=device), weights)  # update model if SourceChangeWarning
# model.fuse()
model.to(device).eval()
if half:
    model.half()  # to FP16


######## Function




#### cosine similarity
def cosine_similarity(a, b):
    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))

#### 画像ファイルから画像の特徴ベクトルを取得(ndarray 512次元)
def feature_vector(image_ndarray):
    img = image_ndarray.copy()
    img = img[:, :, ::-1]
    #img = Image.fromarray(np.uint8(img*255))
    img = Image.fromarray(img)
    #img.save("test.jpg")
    ##img_cropped = mtcnn(img, "test2.jpg")
    if min(img.size) < FaceSize: return None, None, None
    img_cropped, prob = mtcnn(img, return_prob=True)

    if img_cropped is None: return None, None, None
    #img.save("test" + str(prob) + ".jpg")
    feature_vector = resnet(img_cropped.unsqueeze(0))
    feature_vector_np = feature_vector.squeeze().to('cpu').detach().numpy().copy()
    return feature_vector_np, prob, img_cropped

#### 顔画像を保存
tf = transforms.ToPILImage(mode='RGB')
def save_face_img(faceimg, MaxSim, MaxSimID, prob):

    img = (faceimg + 1) / 2 # (-1 ~ 1) to (0 ~ 1)
    img = tf(img) # Tensor to PIL
    savepath = "img/faceid-" + str(MaxSimID) + "_sim-" + str(round(MaxSim,4)) + "_prob-" + str(round(prob,4)) + ".jpg"
    img.save(savepath)

def bbox_rel(image_width, image_height,  *xyxy):
    """" Calculates the relative bounding box from absolute pixel values. """
    bbox_left = min([xyxy[0].item(), xyxy[2].item()])
    bbox_top = min([xyxy[1].item(), xyxy[3].item()])
    bbox_w = abs(xyxy[0].item() - xyxy[2].item())
    bbox_h = abs(xyxy[1].item() - xyxy[3].item())
    x_c = (bbox_left + bbox_w / 2)
    y_c = (bbox_top + bbox_h / 2)
    w = bbox_w
    h = bbox_h
    return x_c, y_c, w, h


def compute_color_for_labels(label):
    """
    Simple function that adds fixed color depending on the class
    """
    color = [int((p * (label ** 2 - label + 1)) % 255) for p in palette]
    return tuple(color)


def draw_boxes(img, bbox, identities=None, offset=(0,0)):
    for i, box in enumerate(bbox):
        x1, y1, x2, y2 = [int(i) for i in box]
        x1 += offset[0]
        x2 += offset[0]
        y1 += offset[1]
        y2 += offset[1]
        # box text and bar
        id = int(identities[i]) if identities is not None else 0    
        #color = compute_color_for_labels(id)
        color = compute_color_for_labels(id + 2)
        if(id == -1): id = label = "?"
        else: label = '{}{:d}'.format("", id)
        t_size = cv2.getTextSize(label, cv2.FONT_HERSHEY_PLAIN, 2 , 2)[0]
        cv2.rectangle(img, (x1, y1),(x2,y2), color, 3)
        cv2.rectangle(img, (x1, y1), (x1 + t_size[0] + 3, y1 + t_size[1] + 4), color, -1)
        cv2.putText(img, label, (x1, y1 + t_size[1] + 4), cv2.FONT_HERSHEY_PLAIN, 2, [255, 255, 255], 2)
    return img


def detect(opt, save_img=False):
    out, source, weights, view_img, save_txt, imgsz = \
        opt.output, opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size
    webcam = source == '0' or '1' or source.startswith('rtsp') or source.startswith('http') or source.endswith('.txt')

    ####
    IDinfoDict = {} # Deepsort ID をKeyとして、class IDInfoを保存する

    # initialize deepsort
    cfg = get_config()
    cfg.merge_from_file(opt.config_deepsort)
    deepsort = DeepSort(cfg.DEEPSORT.REID_CKPT, 
                        max_dist=cfg.DEEPSORT.MAX_DIST, min_confidence=cfg.DEEPSORT.MIN_CONFIDENCE, 
                        nms_max_overlap=cfg.DEEPSORT.NMS_MAX_OVERLAP, max_iou_distance=cfg.DEEPSORT.MAX_IOU_DISTANCE, 
                        max_age=cfg.DEEPSORT.MAX_AGE, n_init=cfg.DEEPSORT.N_INIT, nn_budget=cfg.DEEPSORT.NN_BUDGET, use_cuda=True)

    # Second-stage classifier
    classify = False
    if classify:
        modelc = torch_utils.load_classifier(name='resnet101', n=2)  # initialize
        modelc.load_state_dict(torch.load('weights/resnet101.pt', map_location=device)['model'])  # load weights
        modelc.to(device).eval()

    # Set Dataloader
    vid_path, vid_writer = None, None
    if webcam:
        view_img = True
        cudnn.benchmark = True  # set True to speed up constant image size inference
        dataset = LoadStreams(source, img_size=imgsz)
    else:
        view_img = True
        save_img = True
        dataset = LoadImages(source, img_size=imgsz)

    # Get names and colors
    names = model.module.names if hasattr(model, 'module') else model.names
    colors = [[random.randint(0, 255) for _ in range(3)] for _ in range(len(names))]

    # Run inference
    t0 = time.time()
    img = torch.zeros((1, 3, imgsz, imgsz), device=device)  # init img
    _ = model(img.half() if half else img) if device.type != 'cpu' else None  # run once
    for path, img, im0s, vid_cap in dataset:
        img = torch.from_numpy(img).to(device)
        img = img.half() if half else img.float()  # uint8 to fp16/32
        img /= 255.0  # 0 - 255 to 0.0 - 1.0
        # (640x480) => img=[1, 3, 512, 640] 
        # (640x480) => img=[1, 3, 512, 640] 
        if img.ndimension() == 3:
            img = img.unsqueeze(0)

        # Inference
        t1 = time_synchronized()
        pred = model(img, augment=opt.augment)[0]

        # Apply NMS
        pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, classes=opt.classes, agnostic=opt.agnostic_nms)
        t2 = time_synchronized()

        # Apply Classifier
        if classify:
            pred = apply_classifier(pred, modelc, img, im0s)


        # Process detections
        for i, det in enumerate(pred):  # detections per image
            if webcam:  # batch_size >= 1
                p, s, im0 = path[i], '%g: ' % i, im0s[i].copy()
            else:
                p, s, im0 = path, '', im0s

            save_path = str(Path(out) / Path(p).name)
            txt_path = str(Path(out) / Path(p).stem) + ('_%g' % dataset.frame if dataset.mode == 'video' else '')
            s += '%gx%g ' % img.shape[2:]  # print string
            gn = torch.tensor(im0.shape)[[1, 0, 1, 0]]  #  normalization gain whwh
            if det is not None and len(det):
                # Rescale boxes from img_size to im0 size
                det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()

                # Print results
                for c in det[:, -1].unique():
                    n = (det[:, -1] == c).sum()  # detections per class
                    s += '%g %ss, ' % (n, names[int(c)])  # add to string

                bbox_xywh = []
                confs = []

                # Adapt detections to deep sort input format
                
                for *xyxy, conf, cls in det:
                    img_h, img_w, _ = im0.shape
                    x_c, y_c, bbox_w, bbox_h = bbox_rel(img_w, img_h, *xyxy)
                    obj = [x_c, y_c, bbox_w, bbox_h]
                    bbox_xywh.append(obj)
                    confs.append([conf.item()])
                ####  bbox_xywh[N][x, y, w, h] = DeepSort Bounding Box
                xywhs = torch.Tensor(bbox_xywh)
                confss = torch.Tensor(confs)

                # Pass detections to deepsort
                outputs = deepsort.update(xywhs, confss , im0)

                # draw boxes for visualization
                if len(outputs) > 0:
                    bbox_tlwh = []
                    bbox_xyxy = outputs[:, :4]
                    identities = outputs[:, -1] #### Deepsort IDのリスト
                    #ori_im = draw_boxes(im0, bbox_xyxy, identities) # 画像にバウンディングボックスとID表示

                    #### ↓ 追加 ↓
                    #### DeepSortID（バウンディングボックス）毎にループ
                    for bbox, deepsort_id, conf in zip(bbox_xyxy, identities, confs):
                        if(conf[0] < 0.5): continue # Deepsortの確からしさが低ければ無視して次のDeepSortIDへ
                        # 新しいDeepSortIDが出現したら新しいIDInfoクラスを作成してIDinfoDictに登録する。
                        if(not deepsort_id in IDinfoDict): IDinfoDict[deepsort_id] = IDInfo(deepsort_id)
                            
                        idinfo =  IDinfoDict[deepsort_id] # 各DeepSortIDの情報を保存するクラス
                        idinfo.frameplus() # Nフレームに1回だけ処理をする

                        if(idinfo.frame == 0): # Nフレームに1回だけ処理をする　N=MaxFrame+1
                            print("deepsort_id: ", deepsort_id)
                            print("conf: ", conf[0])
                            #### 画像から顔部分トリミング
                            x1, y1, x2, y2 = bbox
                            cropped = im0[round(y1): round(y2), round(x1): round(x2), :]

                            #### 顔画像から特徴ベクトル抽出
                            fv, prob, faceimg = feature_vector(cropped)
                            if(fv is None):
                                print("None Face")
                                continue # 顔が検出できなければ無視して次のDeepSortIDへ
                            if(prob < 0.80):
                                print("Bad Face, prob=", prob)
                                continue # 顔らしくなければ無視して次のDeepSortIDへ
                            idinfo.addfv(fv, prob) # データベースに特徴ベクトルを追加

                            #### データベースに登録されているIDと特徴ベクトルの組をリストrowsとして取得

                            try:
                                cur.execute("select * from faces")
                                rows = cur.fetchall()
                            except Exception as e:
                                print("\nMySQL ERROR!!!!!!!!\n")
                                print("ERROR",e)
                                continue

                            if(len(rows) == 0): # データベースに登録されているIDがひとつもない場合
                                idinfo.add_newID_to_DB() # 現在のDeepsortIDをデータベースに登録する
                                # save_face_img(faceimg,99, 1, prob)

                            else:
                                MaxSim, MaxSimID = idinfo.calc_sim_with_db(rows)
                                print("MaxSim: ", MaxSim, "MaxSimID: ", MaxSimID, "prob: ", prob)
                                #save_face_img(faceimg, MaxSim, MaxSimID)
                                
                                if(MaxSim <= 0.65):
                                    idinfo.dissim += 1
                                    if(idinfo.dissim >= idinfo.Maxdissim):
                                        idinfo.add_newID_to_DB() # 現在のDeepsortIDをデータベースに登録する
                                    # save_face_img(faceimg, MaxSim, idinfo.face_id, prob)
                                elif(MaxSim >= 0.875):
                                    idinfo.face_id = MaxSimID # データベース上の最も類似度の高い顔IDをidinfoに登録する
                                    # if(prob >= 0.9): idinfo.update_FV_to_DB()
                                    # save_face_img(faceimg, MaxSim, MaxSimID, prob)
                                else:
                                    # save_face_img(faceimg, MaxSim, 99, prob)
                                    pass
                        else:
                            pass

                
                    for i in range(len(identities)):
                        if(identities[i] in IDinfoDict): identities[i] = IDinfoDict[identities[i]].face_id
                        else: identities[i] = -1
                    ori_im = draw_boxes(im0, bbox_xyxy, identities) # 画像にバウンディングボックスとID表示


            # Print time (inference + NMS)
            ########################## print('%sDone. (%.3fs)' % (s, t2 - t1))

            # Stream results
            if view_img:
                cv2.imshow(p, im0)
                #print(p)
                if cv2.waitKey(1) == ord('q'):  # q to quit
                    #### Quit
                    cnx.close()
                    return "Quit"
                    #raise StopIteration

            # Save results (image with detections)
            if save_img:
                if dataset.mode == 'images':
                    cv2.imwrite(save_path, im0)
                else:
                    if vid_path != save_path:  # new video
                        vid_path = save_path
                        if isinstance(vid_writer, cv2.VideoWriter):
                            vid_writer.release()  # release previous video writer

                        fps = vid_cap.get(cv2.CAP_PROP_FPS)
                        w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
                        h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
                        vid_writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*opt.fourcc), fps, (w, h))
                    vid_writer.write(im0)

    if save_txt or save_img:
        print('Results saved to %s' % os.getcwd() + os.sep + out)
        if platform == 'darwin':  # MacOS
            os.system('open ' + save_path)

    print('Done. (%.3fs)' % (time.time() - t0))


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--weights', type=str, default='yolov5/weights/yolov5m.pt', help='model.pt path')
    parser.add_argument('--source', type=str, default='inference/images', help='source')  # file/folder, 0 for webcam
    parser.add_argument('--output', type=str, default='inference/output', help='output folder')  # output folder
    parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)') # default 640
    parser.add_argument('--conf-thres', type=float, default=0.4, help='object confidence threshold')
    parser.add_argument('--iou-thres', type=float, default=0.5, help='IOU threshold for NMS')
    parser.add_argument('--fourcc', type=str, default='mp4v', help='output video codec (verify ffmpeg support)')
    parser.add_argument('--device', default='0', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
    parser.add_argument('--view-img', action='store_true', help='display results')
    parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
    # class 0 is person
    parser.add_argument('--classes', nargs='+', type=int, default=[0], help='filter by class')
    parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')
    parser.add_argument('--augment', action='store_true', help='augmented inference')
    parser.add_argument("--config_deepsort", type=str, default="deep_sort/configs/deep_sort.yaml")
    args = parser.parse_args()
    args.img_size = check_img_size(args.img_size)

    #### カメラ2台使用のため追記：「rtsp://~ はIPカメラ, 0はPCカメラ」
    args2 = copy.deepcopy(args)
    args.source = "0"
    args2.source = "1"
    args.source = "rtsp://dkkenkai:Dk123456@192.168.0.1/MediaInput/h264/stream_3"
    args2.source = "rtsp://dkkenkai:Dk123456@192.168.0.2/MediaInput/h264/stream_3"
    print(args)
    print(args2)

    # Original
    # with torch.no_grad():
    #     detect(args)
    #     detect(args2)


    with torch.no_grad():
        with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:
            camera1 = executor.submit(detect, args)
            camera2 = executor.submit(detect, args2)
            print("end")

        # with concurrent.futures.ProcessPoolExecutor(max_workers=2) as executor:
        #     camera1 = executor.submit(detect, args)
        #     camera2 = executor.submit(detect, args)
        #     print("end")